#include "search_engine/search_index.h"

#include <dirent.h>

#include <set>
#include <boost/lexical_cast.hpp>

#include "mmseg/mmseg.h"
#include "base/logging.h"
#include "search_engine/word_prob.h"

map<double, string> SearchIndex::GetWeightedTokens(const string & text) {
  map<double, string> ret;
  cout << "To search text : "<< text << endl;
  rmmseg::Algorithm * alg = mmseg_algor_create(text.c_str(), text.size());
  Token token = mmseg_next_token(alg);
  while(token.length > 0) {
    string word(token.text, token.length);
    ret.insert(make_pair(word.size() * WordProb::Instance().GetWordProb(word), word));
    token = mmseg_next_token(alg);
  }
  delete alg;
  return ret;;
}

vector<int> SearchIndex::Search(const string & text, int start, size_t count, vector<string> * seg_words, int32_t * total_count) {
  vector<int> res;
  *total_count = 0;
  map<double, string> tokens = GetWeightedTokens(text);
  if (tokens.size() == 0) {
    return res;
  }

  if (tokens.size() == 1) {
    seg_words->push_back(tokens.begin()->second);
    LOG_DEBUG("word : " << tokens.begin()->second);

    boost::shared_lock<boost::shared_mutex> lock(shared_mutex_);
    IndexCache::const_iterator it = index_cache_.find(tokens.begin()->second);
    if (it != index_cache_.end()) {
      const multimap<double, int> & ids = it->second;
      *total_count = ids.size();
      multimap<double, int>::const_reverse_iterator ids_it = ids.rbegin();
      std::advance(ids_it, start);

      for(; ids_it != ids.rend(); ++ids_it) {
        res.push_back(ids_it->second);
        if (res.size() >= count) {
          break;
        }
      }
    }
    return res;
  }
  
  map<double, string>::const_iterator token_it = tokens.begin();
  map<int, double> article_scores;
  for(; token_it != tokens.end(); ++token_it) {
    seg_words->insert(seg_words->begin(), token_it->second);

    boost::shared_lock<boost::shared_mutex> lock(shared_mutex_);
    IndexCache::const_iterator index_it = index_cache_.find(token_it->second);
    if (index_it != index_cache_.end()) {
      const multimap<double, int> & ids = index_it->second;
      multimap<double, int>::const_reverse_iterator ids_it = ids.rbegin();
      for(; ids_it != ids.rend(); ++ids_it) {
        map<int, double>::iterator id_it = article_scores.find(ids_it->second);
        if (id_it == article_scores.end()) {
          article_scores[ids_it->second] = ids_it->first * token_it->first;
        } else {
          id_it->second += ids_it->first * token_it->first;
          id_it->second *= 5.0; // 交集的数据，权重乘以5倍
        }
      }
    }
  }

  *total_count = article_scores.size();

  map<int, double>::iterator score_it = article_scores.begin();
  multimap<double, int> score_articles;
  for(; score_it != article_scores.end(); ++score_it) {
    score_articles.insert(make_pair(score_it->second, score_it->first));
  }

  multimap<double, int>::const_reverse_iterator ids_it = score_articles.rbegin();
  std::advance(ids_it, start);
  for(; ids_it != score_articles.rend(); ++ids_it) {
    res.push_back(ids_it->second);
    if (res.size() >= count) {
      break;
    }
  }
  return res;
}

void SearchIndex::Dump() {
  boost::shared_lock<boost::shared_mutex> lock(shared_mutex_);
  IndexCache::const_iterator i = index_cache_.begin();
  for(; i != index_cache_.end(); ++i) {
    cout << i->first << ':';
    multimap<double, int>::const_iterator j = i->second.begin();
    for(; j != i->second.end(); ++j) {
      cout << j->first << '-' << j->second << ",";
    }
    cout << endl;
  }
}

vector<int> SearchIndex::SearchByTime(const string & text, int start, size_t count, vector<string> * seg_words, int32_t * total_count) {
  vector<int> res;
  *total_count = 0;
  map<double, string> tokens = GetWeightedTokens(text);
  if (tokens.size() == 0) {
    return res;
  }

  if (tokens.size() == 1) {
    seg_words->push_back(tokens.begin()->second);
    LOG_DEBUG("word : " << tokens.begin()->second);
    boost::shared_lock<boost::shared_mutex> lock(shared_mutex_);
    IndexCache::const_iterator it = index_cache_.find(tokens.begin()->second);
    if (it != index_cache_.end()) {
      const multimap<double, int> & ids = it->second;
      *total_count = ids.size();

      multimap<double, int>::const_reverse_iterator ids_it = ids.rbegin();

      set<int> sort_ids;
      for(; ids_it != ids.rend(); ++ids_it) {
        sort_ids.insert(ids_it->second);
      }

      for(set<int>::const_iterator j = sort_ids.begin(); j != sort_ids.end(); ++j) {
        LOG_DEBUG("search id " << *j);
      }

      set<int>::const_reverse_iterator sort_ids_it = sort_ids.rbegin();
      std::advance(sort_ids_it, start);

      for(; sort_ids_it != sort_ids.rend(); ++sort_ids_it) {
        res.push_back(*sort_ids_it);
        LOG_DEBUG("return id " << *sort_ids_it);
        if (res.size() >= count) {
          break;
        }
      }
    }
    return res;
  }

  map<double, string>::const_iterator token_it = tokens.begin();
  map<int, double> article_scores;
  for(; token_it != tokens.end(); ++token_it) {
    seg_words->insert(seg_words->begin(), token_it->second);

    boost::shared_lock<boost::shared_mutex> lock(shared_mutex_);
    IndexCache::const_iterator index_it = index_cache_.find(token_it->second);
    if (index_it != index_cache_.end()) {
      const multimap<double, int> & ids = index_it->second;
      multimap<double, int>::const_reverse_iterator ids_it = ids.rbegin();
      for(; ids_it != ids.rend(); ++ids_it) {
        map<int, double>::iterator id_it = article_scores.find(ids_it->second);
        if (id_it == article_scores.end()) {
          article_scores[ids_it->second] = ids_it->first * token_it->first;
        } else {
          id_it->second += ids_it->first * token_it->first;
          id_it->second *= 5.0; // 交集的数据，权重乘以5倍
        }
      }
    }
  }

  *total_count = article_scores.size();

  map<int, double>::iterator score_it = article_scores.begin();
  multimap<double, int> score_articles;
  for(; score_it != article_scores.end(); ++score_it) {
    score_articles.insert(make_pair(score_it->second, score_it->first));
  }

  multimap<double, int>::const_reverse_iterator ids_it = score_articles.rbegin();

  set<int> sort_ids;
  for(; ids_it != score_articles.rend(); ++ids_it) {
    sort_ids.insert(ids_it->second);
  }

  for(set<int>::const_iterator i = sort_ids.begin(); i != sort_ids.end(); ++i) {
    LOG_DEBUG("search id " << *i);
  }

  set<int>::const_reverse_iterator sort_ids_it = sort_ids.rbegin();
  std::advance(sort_ids_it, start);

  for(; sort_ids_it != sort_ids.rend(); ++sort_ids_it) {
    res.push_back(*sort_ids_it);
    LOG_DEBUG("return id " << *sort_ids_it);
    if (res.size() >= count) {
      break;
    }
  }

  return res;
}

bool SearchIndex::LoadIndexFile(const char * index_file) {
  FILE * data = fopen(index_file, "r");
  if (!data) {
    LOG_WARN("open data file " << index_file << " error");
    cout << "open data file " << index_file << " error" << endl;
    return false;
  }
  LOG_INFO("Start loading index file " << index_file << " ...");
  cout << "Start loading index file " << index_file << " ..." << endl;
  static const int kMaxLine = 128 * 1024;
  char * line = new char[kMaxLine];
  line[kMaxLine - 1] = 0;
  std::vector<std::string> strs;
  
  while(fgets(line, kMaxLine - 1, data)) {
    boost::split(strs, line, boost::is_any_of(" \r\n"), boost::token_compress_on);
    if (strs.size() != 4) {
      continue;
    }
    int aid = boost::lexical_cast<int>(strs[0]);
    string word(strs[1]);
    boost::trim(word);
    double score = boost::lexical_cast<double>(strs[2]);

    // cout << "aid=" << aid << " word=" << word << word.size() << " score=" << score << endl;
    if (word.size() > 3) {
      boost::unique_lock<boost::shared_mutex> lock(shared_mutex_);
      index_cache_[word].insert(make_pair(score, aid));
    }
  }
  LOG_INFO("Finish loading index file " << index_file << ".");
  delete [] line;
  fclose(data);
  return true;
}

bool SearchIndex::Reload(const char * index_dir) {
  WordProb::Instance().Load();
  if (mmseg_load_chars("./chars.dic") == 0) {
    cerr << "load chars error!" << endl;
  }
  LOG_DEBUG("load chars OK!");

  if (mmseg_load_words("./words.dic") == 0) {
    cerr << "load words error!" << endl;
  }
  LOG_DEBUG("load words OK!");

  DIR* dp = opendir(index_dir);
  struct dirent* de;
  
  int count = 0;
  while (de = readdir(dp)) { 
    if (de->d_type != DT_REG) // Only print regular files
       continue;
     cout << "Found file " << ++count << " " << string(index_dir) + de->d_name << endl;
     LoadIndexFile((string(index_dir) + de->d_name).c_str());
  }
  closedir(dp);
  return true;
}


